# -*- coding:utf-8 -*-
from fastapi import FastAPI, Request, WebSocket, Depends
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.responses import HTMLResponse
import uvicorn
import asyncio

from starlette.websockets import WebSocketState

app = FastAPI()  # 创建 api 对象


class ReqeustStats:
    def __init__(self, ts, tps, succ_req, err_req, abort_req, retry_req, unknown_req, fault_id):
        self.ts = ts
        self.tps = tps
        self.succ_req = succ_req
        self.err_req = err_req
        self.abort_req = abort_req
        self.retry_req = retry_req
        self.unknown_req = unknown_req
        self.fault_id = fault_id

    def __str__(self):
        return ("ts:%s,tps:%s, succ:%s, err:%s, ... , fault:%s" %
                (self.ts, self.tps, self.succ_req, self.err_req, self.fault_id))


@app.get("/")  # 根路由
def root():
    return {"Hello": "DataBase Reliability benchmark"}


@app.get("/say/{data}")
def say(data: str, q: int):
    return {"data": data, "item": q}


app.mount("/static", StaticFiles(directory="static"), name="static")  # 挂载静态文件，指定目录

templates = Jinja2Templates(directory="templates")  # 模板目录

txn_stats = []
keep_size = 20
pre_succ_req = 0
pre_err_req = 0
pre_abort_req = 0
pre_retry_req = 0
pre_unknown_req = 0


@app.get("/txn")
async def sync_req_stats(request: Request):
    global txn_stats, keep_size
    return templates.TemplateResponse(
        "index.html",
        {"request": request, "txn_stats": txn_stats[:keep_size],
            # "txn/sec": txn_stats[-1].tps, # TypeError: 'ReqeustStats' object is not subscriptable
            "pre_succ":pre_succ_req,
            "pre_err":pre_err_req,
            "pre_abort":pre_abort_req,
            "pre_retry":pre_retry_req,
            "pre_unknown":pre_unknown_req
         }
    )


html = """
<!DOCTYPE html>
<html>
    <head>
        <title>Chat</title>
    </head>
    <body>
        <h1>WebSocket Chat</h1>
        <form action="" onsubmit="sendMessage(event)">
            <input type="text" id="messageText" autocomplete="off"/>
            <button>Send</button>
        </form>
        <ul id='messages'>
        </ul>
        <script>
            var ws = new WebSocket("ws://localhost:8000/ws_html_client");
            ws.onmessage = function(event) {
                var messages = document.getElementById('messages')
                var message = document.createElement('li')
                var content = document.createTextNode(event.data)
                message.appendChild(content)
                messages.appendChild(message)
            };
            function sendMessage(event) {
                var input = document.getElementById("messageText")
                ws.send(input.value)
                input.value = ''
                event.preventDefault()
            }
        </script>
    </body>
</html>
"""


@app.get("/wsget")
async def wsget():
    return HTMLResponse(html)


@app.websocket("/ws_html_client")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        print('ws_html_client.receive_text:', data)
        await websocket.send_text(f"Message text was: {data}")


@app.websocket("/ws_java_client")
async def websocket_endpoint(websocket: WebSocket):
    global txn_stats, keep_size, pre_succ_req, pre_err_req, pre_abort_req, pre_retry_req, pre_unknown_req
    await websocket.accept()
    while True:
        try:
            if websocket.client_state != WebSocketState.CONNECTED:
                print("Not CONNECTED")
                break

            trxs = await websocket.receive_text()
            trx_nums = trxs.split(",")

            tps_ = float(trx_nums[1])
            succ_ = int(trx_nums[2])
            err_ = int(trx_nums[3])
            abort_ = int(trx_nums[4])
            retry_ = int(trx_nums[5])
            unknown_ = int(trx_nums[6])

            req_stats = ReqeustStats(ts=trx_nums[0],
                                     tps=tps_,
                                     succ_req=succ_ - pre_succ_req,
                                     err_req=err_ - pre_err_req,
                                     abort_req=abort_ - pre_abort_req,
                                     retry_req=retry_ - pre_retry_req,
                                     unknown_req=unknown_ - pre_unknown_req,
                                     fault_id=float(trx_nums[7]))


            txn_stats.append(req_stats)

            pre_succ_req = succ_
            pre_err_req = err_
            pre_abort_req = abort_
            pre_retry_req = retry_
            pre_unknown_req = unknown_

            if len(txn_stats) > keep_size:
                txn_stats.pop(0)

            print("work_stats: ", req_stats)
            # await websocket.send_text(f"WS server got data size: [ {len(work_stats)}]")

            # await asyncio.sleep(1)

        except BaseException as e:
            print(e)
            break


if __name__ == '__main__':
    uvicorn.run(app, host="127.0.0.1", port=8000)
